package com.shujia.spark.mllib

import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressionModel}
import org.apache.spark.ml.linalg.{SparseVector, Vectors}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object Demo04ImageModelTrain {
  def main(args: Array[String]): Unit = {

    // 手写数字识别
    // 使用模型
    val spark: SparkSession = SparkSession
      .builder()
      .appName(this.getClass.getSimpleName.replace("$", ""))
      .master("local[*]") // 设置运行的方式
      .config("spark.sql.shuffle.partitions", "16")
      .getOrCreate()

    val imageDF: DataFrame = spark
      .read
      .format("image")
      //      .load("spark/data/mllib/data/test/img/")
      .load("C:\\Users\\qx\\Desktop\\train")

    val imageResDF: DataFrame = spark
      .read
      .format("csv")
      .option("sep", " ")
      .schema("fileName String,label Double")
      .load("spark/data/mllib/data/image_res.txt")

    imageDF.printSchema()

    import spark.implicits._
    val trainDF: DataFrame = imageDF
      .select($"image.origin", $"image.data")
      .as[(String, Array[Byte])]
      .rdd
      .map(t2 => {
        val fileName: String = t2._1.split("/").reverse.head
        val features: SparseVector = Vectors.dense(t2._2.map(b => {
          if (b >= 0 && b < 16)
            0.0
          else
            255.0
        })).toSparse
        (fileName, features)
      })
      .toDF("fileName", "features")
      .join(imageResDF, "fileName")


    val lr: LogisticRegression = new LogisticRegression()
      .setMaxIter(10)
//      .setRegParam(0.3)
//      .setElasticNetParam(0.8)
      .setFitIntercept(true)
      .setFamily("multinomial")

    val imageModel: LogisticRegressionModel = lr.fit(trainDF)

    val testDF: DataFrame = trainDF.sample(false, 0.001)

    val testResDF: DataFrame = imageModel.transform(testDF)
    testResDF.cache()

    testResDF.show()

    println(s"模型的准确率为：${testResDF.where($"label" === $"prediction").count().toDouble / testResDF.count()}")

    testResDF.unpersist()

  }

}
